package szte.objectorgy.collections;

/**
 * Map data structure with a left-leaning red-black tree underlying
 * implementation.
 * 
 * @author Endre PALATINUS
 * @param <K>
 *            The type of the key used in the map.
 * @param <V>
 *            The type of the values stored in the map.
 * 
 */
public class TreeMap<K extends Comparable<K>, V extends Comparable<V>>
		implements Map<K, V> {

	private static final boolean RED = true;
	private static final boolean BLACK = false;

	protected RedBlackTreeNode root;

	/**
	 * A node of the search tree.
	 */
	public class RedBlackTreeNode implements Comparable<RedBlackTreeNode> {

		public K key;
		public V value;

		public RedBlackTreeNode left, right;

		public boolean color;

		/** Number of RedBlackTreeNodes reachable from this node */
		private int subTreeSize;

		public RedBlackTreeNode(K key, V value) {
			this.color = RED;
			this.subTreeSize = 1;

			this.key = key;
			this.value = value;					
		}

		@Override
		public int compareTo(RedBlackTreeNode arg0) {
			return key.compareTo(arg0.key);
		}
		
		public String toString() {
			return key.toString() + ": " + value.toString();
		}
		
		public void accept(TreeVisitor v) {
			v.visit(this);			
		}

		public K getKey() {
			return key;
		}

		public V getValue() {
			return value;
		}

		public RedBlackTreeNode getLeft() {
			return left;
		}

		public RedBlackTreeNode getRight() {
			return right;
		}

		public boolean isColor() {
			return color;
		}
	}

	/**
	 * Iterator for the inorder tree traversal of a TreeMap.
	 * 
	 * @author Endre PALATINUS
	 */
	protected class TreeMapIterator implements Iterator<Pair<K, V>> {

		private TreeMap<K, V>.RedBlackTreeNode current;
		private TreeMap<K, V>.RedBlackTreeNode next;

		private Stack<TreeMap<K, V>.RedBlackTreeNode> stack;

		@SuppressWarnings("unchecked")
		public TreeMapIterator() {
			current = null;
			next = root;

			stack = new Stack<TreeMap<K, V>.RedBlackTreeNode>(
					(Class<RedBlackTreeNode>) root.getClass());

			while (next.left != null) {
				stack.push(next);
				next = next.left;
			}
		}

		@Override
		public boolean hasNext() {
			return next != null;
		}

		@Override
		public Pair<K, V> next() {
			current = next;
			Pair<K, V> result = new Pair<K, V>(current.key, current.value);

			if (next.right == null && !stack.isEmpty()) {
				next = stack.pop();
			} else {
				next = next.right;

				if (next != null) {
					while (next.left != null) {
						stack.push(next);
						next = next.left;
					}
				}
			}

			return result;
		}
	}

	protected TreeMapIterator getIterator() {
		return new TreeMapIterator();
	}

	public int size() {
		return size(root);
	}

	private int size(RedBlackTreeNode x) {
		if (x == null)
			return 0;
		else
			return x.subTreeSize;
	}

	public V get(K key) {
		return get(root, key);
	}

	private V get(RedBlackTreeNode x, K key) {
		if (x == null)
			return null;
		if (eq(key, x.key))
			return x.value;
		if (less(key, x.key))
			return get(x.left, key);
		else
			return get(x.right, key);
	}

	protected K min(RedBlackTreeNode x) {
		if (x.left == null)
			return x.key;
		else
			return min(x.left);
	}

	protected K max(RedBlackTreeNode x) {
		if (x.right == null)
			return x.key;
		else
			return max(x.right);
	}

	public void put(K key, V value) {
		root = insert(root, key, value);
		root.color = BLACK;
	}

	private RedBlackTreeNode insert(RedBlackTreeNode h, K key, V value) {
		if (h == null)
			return new RedBlackTreeNode(key, value);

		if (isRed(h.left) && isRed(h.right))
			colorFlip(h);

		if (eq(key, h.key))
			h.value = value;
		else if (less(key, h.key))
			h.left = insert(h.left, key, value);
		else
			h.right = insert(h.right, key, value);

		if (isRed(h.right))
			h = rotateLeft(h);

		if (isRed(h.left) && isRed(h.left.left))
			h = rotateRight(h);

		return fixSubTreeSize(h);
	}

	private RedBlackTreeNode deleteMin(RedBlackTreeNode h) {
		if (h.left == null)
			return null;

		if (!isRed(h.left) && !isRed(h.left.left))
			h = moveRedLeft(h);

		h.left = deleteMin(h.left);

		return ensureLeftLeaningProperty(h);
	}

	private RedBlackTreeNode deleteMax(RedBlackTreeNode h) {

		if (isRed(h.left))
			h = rotateRight(h);

		if (h.right == null)
			return null;

		if (!isRed(h.right) && !isRed(h.right.left))
			h = moveRedRight(h);

		h.right = deleteMax(h.right);

		return ensureLeftLeaningProperty(h);
	}

	private RedBlackTreeNode delete(RedBlackTreeNode h, K key) {
		if (less(key, h.key)) {
			if (!isRed(h.left) && !isRed(h.left.left))
				h = moveRedLeft(h);
			h.left = delete(h.left, key);
		} else {
			if (isRed(h.left))
				h = rotateRight(h);
			if (eq(key, h.key) && (h.right == null))
				return null;
			if (!isRed(h.right) && !isRed(h.right.left))
				h = moveRedRight(h);
			if (eq(key, h.key)) {
				h.value = get(h.right, min(h.right));
				h.key = min(h.right);
				h.right = deleteMin(h.right);
			} else
				h.right = delete(h.right, key);
		}

		return ensureLeftLeaningProperty(h);
	}

	// Helper methods

	private boolean less(K a, K b) {
		return a.compareTo(b) < 0;
	}

	private boolean eq(K a, K b) {
		return a.compareTo(b) == 0;
	}

	private boolean isRed(RedBlackTreeNode x) {
		if (x == null)
			return false;
		return (x.color == RED);
	}

	private void colorFlip(RedBlackTreeNode h) {
		h.color = !h.color;
		h.left.color = !h.left.color;
		h.right.color = !h.right.color;
	}

	private RedBlackTreeNode rotateLeft(RedBlackTreeNode h) {
		RedBlackTreeNode x = h.right;
		h.right = x.left;
		x.left = fixSubTreeSize(h);
		x.color = x.left.color;
		x.left.color = RED;
		return fixSubTreeSize(x);
	}

	private RedBlackTreeNode rotateRight(RedBlackTreeNode h) {
		RedBlackTreeNode x = h.left;
		h.left = x.right;
		x.right = fixSubTreeSize(h);
		x.color = x.right.color;
		x.right.color = RED;
		return fixSubTreeSize(x);
	}

	private RedBlackTreeNode moveRedLeft(RedBlackTreeNode h) {
		colorFlip(h);
		if (isRed(h.right.left)) {
			h.right = rotateRight(h.right);
			h = rotateLeft(h);
			colorFlip(h);
		}
		return h;
	}

	private RedBlackTreeNode moveRedRight(RedBlackTreeNode h) {
		colorFlip(h);
		if (isRed(h.left.left)) {
			h = rotateRight(h);
			colorFlip(h);
		}
		return h;
	}

	private RedBlackTreeNode ensureLeftLeaningProperty(RedBlackTreeNode h) {
		if (isRed(h.right))
			h = rotateLeft(h);

		if (isRed(h.left) && isRed(h.left.left))
			h = rotateRight(h);

		if (isRed(h.left) && isRed(h.right))
			colorFlip(h);

		return fixSubTreeSize(h);
	}

	private RedBlackTreeNode fixSubTreeSize(RedBlackTreeNode h) {
		h.subTreeSize = size(h.left) + size(h.right) + 1;
		return h;
	}

	public String toString() {
		return toString(root);
	}

	private String toString(RedBlackTreeNode x) {
		StringBuilder sb = new StringBuilder();
		
		TreeMapIterator it = getIterator();
		while (it.hasNext()) {
			sb.append(it.next() + "; ");
		}
		
		return sb.toString();
	}

	@Override
	public void clear() {
		root = null;
	}

	@Override
	public boolean containsKey(K key) {
		return (get(key) != null);
	}

	@Override
	public boolean containsValue(V value) {
		// TODO Auto-generated method stub
		return false;
	}

	@Override
	public Set<Pair<K, V>> entrySet() {
		Set<Pair<K, V>> result = new TreeSet<Pair<K, V>>();

		TreeMapIterator it = new TreeMapIterator();
		while (it.hasNext()) {
			result.add(it.next());
		}

		return result;
	}

	@Override
	public boolean isEmpty() {
		return root == null;
	}

	@Override
	public Set<K> keySet() {
		Set<K> result = new TreeSet<K>();

		TreeMapIterator it = new TreeMapIterator();
		while (it.hasNext()) {
			result.add(it.next().getKey());
		}

		return result;
	}

	@Override
	public void putAll(Map<K, V> map) {
		Iterator<Pair<K, V>> it = map.entrySet().getIterator();
		while (it.hasNext()) {
			Pair<K, V> entry = it.next();
			put(entry.getKey(), entry.getValue());
		}
	}

	@Override
	public void remove(K key) {
		root = delete(root, key);
		root.color = BLACK;
	}

	@Override
	public Collection<V> values() {
		@SuppressWarnings("unchecked")
		Vector<V> result = new Vector<V>((Class<V>) root.value.getClass());
		TreeMapIterator it = new TreeMapIterator();

		while (it.hasNext()) {
			result.add(it.next().getValue());
		}

		return result;
	}

}
